#include "tcp_sender.hh"

#include "tcp_config.hh"

#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&... /* unused */) {}

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _stream(capacity)
    , _time_to_retransmit(0)
    , _retransmit_timeout(_initial_retransmission_timeout)
    , _consecutive_retransmissions(0)
    , _next_seqno(0)
    , _unack_size(0) {}

uint64_t TCPSender::bytes_in_flight() const { return _unack_size; }

void TCPSender::fill_window() {
    if (!_syn_set) {
        TCPSegment segment;
        _syn_set = true;
        segment.header().syn = true;
        send_segment(segment);
    }

    while (_stream.buffer_size() && (bytes_in_flight() < _received_window_size)) {
        size_t length = min(_received_window_size - bytes_in_flight(), TCPConfig::MAX_PAYLOAD_SIZE);
        TCPSegment segment;
        segment.payload() = Buffer(std::move(_stream.read(length)));
        if (!_fin_set && _stream.eof() &&
            _received_window_size - segment.length_in_sequence_space() - bytes_in_flight() > 0) {
            _fin_set = true;
            segment.header().fin = true;
        }
        if (segment.length_in_sequence_space() == 0)
            return;
        send_segment(segment);
    }

    if (!_fin_set && _stream.eof() && (!bytes_in_flight() || bytes_in_flight() < _received_window_size)) {
        TCPSegment segment;
        segment.header().fin = true;
        _fin_set = true;
        send_segment(segment);
    }
}

void TCPSender::send_segment(TCPSegment &seg) {
    seg.header().seqno = next_seqno();
    _segments_out.push(seg);
    _outstanding.push(seg);
    _unack_size += seg.length_in_sequence_space();
    _next_seqno += seg.length_in_sequence_space();
}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
//! \returns `false` if the ackno appears invalid (acknowledges something the TCPSender hasn't sent yet)
bool TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    uint64_t received_ackno = unwrap(ackno, _isn, _next_seqno);
    if (received_ackno > _next_seqno) {
        return false;
    }
    _received_window_size = window_size;
    _consecutive_retransmissions = 0;
    _retransmit_timeout = _initial_retransmission_timeout;
    while (!_outstanding.empty()) {
        TCPSegment seg = _outstanding.front();
        if (unwrap(seg.header().seqno, _isn, _next_seqno) + seg.length_in_sequence_space() <= received_ackno) {
            _outstanding.pop();
            _unack_size -= seg.length_in_sequence_space();
            _time_to_retransmit = 0;
        } else {
            break;
        }
    }
    fill_window();
    return true;
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    _time_to_retransmit += ms_since_last_tick;
    if (_time_to_retransmit >= _retransmit_timeout) {
        if (!_outstanding.empty()) {
            _segments_out.push(_outstanding.front());
            _retransmit_timeout *= 2;
            ++_consecutive_retransmissions;
        }
        _time_to_retransmit = 0;
    }
}

unsigned int TCPSender::consecutive_retransmissions() const { return _consecutive_retransmissions; }

void TCPSender::send_empty_segment() {
    TCPSegment segment;
    segment.header().seqno = next_seqno();
    _segments_out.push(segment);
}
